#!/usr/bin/python3
# -*- coding: utf-8 -*-
import os

import sys
# sys.path.append('/opt/work/caffe/python')
sys.path.insert(0, '.')

import cv2
import numpy as np
import struct

############################################################################

# base_path = r'/rootfs/media/kasim/DataSet/Record'
# video_list_path = r'/rootfs/media/kasim/DataSet/Record'
# video_list_bbox_count_file = r'face_video_list_bbox_count.txt'
# video_list_bbox_file = r'face_video_list_bbox.dat'

# base_path = r'/rootfs/media/kasim/DataSet/FaceVideo'
# video_list_path = r'/rootfs/media/kasim/DataSet/FaceVideo'
# # video_list_bbox_count_file = r'video_list_bbox_count.txt'
# # video_list_bbox_file = r'video_list_bbox.dat'
# video_list_bbox_count_file = r'face_video_list_bbox_count.txt'
# video_list_bbox_file = r'face_video_list_bbox.dat'

base_path = r'/rootfs/media/kasim/DataSet/FaceVideo'
video_list_path = r'/rootfs/media/kasim/DataSet/FaceVideo/CropImageMargin'
video_list_bbox_count_file = r'select_bbox_list.txt'
video_list_bbox_file = r'select_bbox.dat'

SLEEP_TIME = 40
# SLEEP_TIME = 0

DISPLAY_WIDTH = 1280
DISPLAY_HEIGHT = 720

SHOW = True
THRESHOLDS = [
    0.5,  # Person
    0.7,  # Cat
    0.7,  # Dog
    0.5,  # BabyCar
    0.3,  # Face
]

# SHOW = False
# THRESHOLDS = [
#     0.05,  # Person
#     0.05,  # Cat
#     0.05,  # Dog
#     0.05,  # BabyCar
#     0.05,  # Face
# ]

COLORS = [
    (255, 0, 0),
    (0, 255, 0),
    (0, 255, 255),
    (255, 0, 0),
    (0, 0, 255),
]

CLASS_NAMES = [
    'Person',
    'Cat',
    'Dog',
    'BabyCar',
    'Face',
]

frame_count = 0
file_count = 0
is_break = False
video_list = []
with open(os.path.join(video_list_path, video_list_bbox_count_file), 'r') as file:
    for video_info in file.readlines():
        video_info = video_info.split(',')
        video_list.append([video_info[0].strip(), int(video_info[1].strip())])

with open(os.path.join(video_list_path, video_list_bbox_file), 'rb') as file:
    for video_name, bbox_count in video_list:
        bbox_info_dict = {}
        for j in range(bbox_count):
            bbox_info = file.read(28)
            if len(bbox_info) < 28:
                is_break = False
                break
            bbox_info = struct.unpack('6i1f', bbox_info)
            frame_id = bbox_info[0]
            label = bbox_info[1]
            score = bbox_info[6]
            if score >= THRESHOLDS[label]:
                if frame_id not in bbox_info_dict.keys():
                    bbox_info_dict[frame_id] = [bbox_info[1:]]
                else:
                    bbox_info_dict[frame_id].append(bbox_info[1:])

        video_path = os.path.join(base_path, video_name)
        file_count += 1
        print('File Count:', file_count, ' Frame Count:', frame_count, video_name)
        cap = cv2.VideoCapture(video_path)

        frame_id = 0
        bbox_count = 0
        while True:
            grabbed, image = cap.read()

            if not grabbed:
                break
            frame_count += 1
            if 0 == (frame_count % 100):
                print('Frame Count:', frame_count)

            bbox_info_list = bbox_info_dict.get(frame_id, None)
            if bbox_info_list is not None:
                for bbox_info in bbox_info_list:
                    lable = bbox_info[0]
                    color = COLORS[lable]
                    class_name = CLASS_NAMES[lable]
                    bbox = bbox_info[1:]
                    label_text = '{} {:.2f}'.format(class_name, bbox[4])
                    cv2.putText(image, label_text, (bbox[0], bbox[1]+32), cv2.FONT_HERSHEY_COMPLEX, 1.0, color)
                    cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, thickness=2)
            frame_id += 1

            if (image.shape[0] != DISPLAY_HEIGHT) or (image.shape[1] != DISPLAY_WIDTH):
                image = cv2.resize(image, (DISPLAY_WIDTH, DISPLAY_HEIGHT))

            cv2.imshow('image', image)
            if 0 == SLEEP_TIME:
                k = cv2.waitKey()
            else:
                k = cv2.waitKey(SLEEP_TIME)
            if k == 27:  # Esc key to stop
                is_break = True
                break
        if is_break:
            break

print('Frame Count:', frame_count)
print('Finish!')


